# -*- coding: utf-8 -*-

import os
import copy
import shutil
import oyaml as yaml


def read_yaml_file(yaml_file_path):
    with open(yaml_file_path, "r", encoding="utf-8") as yaml_file:
        data = yaml.load(yaml_file)
    return data


def write_yaml_file(data, yaml_file_path):
    with open(yaml_file_path, "w", encoding="utf-8") as yaml_file:
        yaml.dump(
            data, yaml_file, allow_unicode=True, default_flow_style=False)


def copy_pair(src_root_folder, dst_root_folder, pair_name, pair):
    pair_index = pair_name.split("_")[1]
    os.makedirs(
        os.path.join(dst_root_folder, pair_index, "gallery"), exist_ok=True)
    os.makedirs(
        os.path.join(dst_root_folder, pair_index, "query"), exist_ok=True)
    # gallery
    src_gallery_image_folder = os.path.join(src_root_folder,
                                            pair["gallery"]["image"])
    dst_gallery_image_folder = os.path.join(dst_root_folder, pair_index,
                                            "gallery", "image")
    if os.path.isdir(dst_gallery_image_folder):
        shutil.rmtree(dst_gallery_image_folder)
    shutil.copytree(src_gallery_image_folder, dst_gallery_image_folder)
    src_gallery_reid_feature = os.path.join(src_root_folder,
                                            pair["gallery"]["reid_feature"])
    dst_gallery_reid_feature = os.path.join(dst_root_folder, pair_index,
                                            "gallery")
    shutil.copy2(src_gallery_reid_feature,
                 dst_gallery_reid_feature)  # 拷贝到gallery文件夹
    # query
    src_query_image_folder = os.path.join(src_root_folder,
                                          pair["query"]["image"])
    dst_query_image_folder = os.path.join(dst_root_folder, pair_index, "query",
                                          "image")
    if os.path.isdir(dst_query_image_folder):
        shutil.rmtree(dst_query_image_folder)
    shutil.copytree(src_query_image_folder, dst_query_image_folder)
    src_query_reid_feature = os.path.join(src_root_folder,
                                          pair["query"]["reid_feature"])
    dst_query_reid_feature = os.path.join(dst_root_folder, pair_index, "query")
    shutil.copy2(src_query_reid_feature, dst_query_reid_feature)  # 拷贝到query文件夹

    # remove dst_root_folder prefix
    return (os.path.join(pair_index, "gallery", "image"),
            os.path.join(pair_index, "gallery",
                         os.path.basename(src_gallery_reid_feature)),
            os.path.join(pair_index, "query", "image"),
            os.path.join(pair_index, "query",
                         os.path.basename(src_query_reid_feature)))


def main():
    src_yaml_file_path = r"G:\work\跨线OD\log\ground_truth.yaml"
    src_data = read_yaml_file(src_yaml_file_path)
    src_root_folder = src_data["root"]

    dst_data = copy.copy(src_data)
    dst_root_folder = r"G:\work\跨线OD\match_data"
    dst_data["root"] = dst_root_folder

    pairs = src_data["pairs"]
    for pair_name in pairs:
        gallery_image, gallery_reid, query_image, query_reid = copy_pair(
            src_root_folder, dst_root_folder, pair_name, pairs[pair_name])
        dst_data["pairs"][pair_name]["gallery"]["image"] = gallery_image
        dst_data["pairs"][pair_name]["gallery"]["reid_feature"] = gallery_reid
        dst_data["pairs"][pair_name]["query"]["image"] = query_image
        dst_data["pairs"][pair_name]["query"]["reid_feature"] = query_reid
    dst_yaml_file_path = os.path.join(dst_root_folder, "match.yaml")
    write_yaml_file(dst_data, dst_yaml_file_path)


if __name__ == "__main__":
    main()
